package com.example.calcite;

import com.zaxxer.hikari.HikariConfig;
import com.zaxxer.hikari.HikariDataSource;
import org.apache.calcite.jdbc.CalciteConnection;
import org.apache.calcite.schema.Schema;
import org.apache.calcite.schema.SchemaPlus;

import java.sql.Connection;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.Map;
import java.util.WeakHashMap;

public class CalciteDataSource extends HikariDataSource {
    private final Map<String, Schema> schemas = new HashMap<>();
    private final Map<Connection, Boolean> registeredConnections = new WeakHashMap<>();

    public CalciteDataSource() {
        super();
    }

    public CalciteDataSource(HikariConfig configuration) {
        super(configuration);
    }

    public void registerSchema(String name, Schema schema) {
        schemas.put(name, schema);
        // 清除已注册连接的记录，因为 schema 发生了变化
        registeredConnections.clear();
    }

    @Override
    public Connection getConnection() throws SQLException {
        Connection connection = super.getConnection();
        
        // 检查这个连接是否已经注册过 schema
        if (!registeredConnections.containsKey(connection)) {
            CalciteConnection calciteConnection = connection.unwrap(CalciteConnection.class);
            SchemaPlus rootSchema = calciteConnection.getRootSchema();
            
            // Register schemas
            for (Map.Entry<String, Schema> entry : schemas.entrySet()) {
                rootSchema.add(entry.getKey(), entry.getValue());
            }
            
            // 标记这个连接已经注册过 schema
            registeredConnections.put(connection, true);
        }
        
        return connection;
    }
} 